"""------------------------------------------------------------------------
 Adiciona imagens do Google Static Maps ao fundo do ArcMap.

 Atente-se aos Termos de Serviço do Google antes de usar as
 imagens para fins comerciais. Maiores informações em:
 https://developers.google.com/maps/documentation/staticmaps

 Google and Bing uses different input and an output spatial reference.
 So, it is necessary to give google maps a WGS84 (epsg:4326) lat/lon
 and it will give back a tile in WebMercator spherical ellipsoid
 (epsg:3857). 

 Criado por Eng. Demian Andrade
 Brasil. Abril, 2014.
------------------------------------------------------------------------"""

"""
BUGS & Ajustes necessários
- ajustar para latitudes/longitudes positivas. (decidir o sinal)
- needs warping to distort staticmap from webmercator to wgs84
- range[0] doesn't work when there's a tree 
"""

__version__ = "$Version: 1.1 $"
__author__ = "Eng. Demian Andrade"

import arcpy
import os
import urllib2
import shutil
import random
from random import randrange


class GlobalMercator(object):
    "Use: GlobalMercator(tilesiez)"
    
    def __init__(self, tileSize):
        "Initialize TMS Global Mercator"

        self.tileSize = tileSize
        self.initialResolution = 2*6378137*math.pi / self.tileSize
        self.originShift = 2*6378137*math.pi / 2.0

    def domain(self, lat, lon):
        "verify domain lat [-90,90] & lon [-180,180]"

        if lat > 90.0:
            error_msg = "Latitude > 90."
        elif lat < -90.0:
            error_msg = "Latitude < -90."
        elif lon > 180.0:
            error_msg = "Longitude > 180."
        elif lon < -180.0:
            error_msg = "Longitude < -180."
        else: error_msg = None

        return error_msg

    def zoomScale(self, scale):
        "Define near zoom from map scale and returns both"
        
        LIST_ZOOM = [21, 20, 19, 18, 17, 16, 15, 14, 13, 12,
                    11, 10, 9, 8, 7 , 6, 5, 4, 3, 2, 1]
        LIST_SCALES = [564.24861, 1128.497220, 2256.994440, 4513.988880,
                      9027.977761, 18055.955520, 36111.911040,
                      72223.822090, 144447.644200, 288895.288400,
                      577790.576700, 1155581.153000, 2311162.307000,
                      4622324.614000, 9244649.227000, 18489298.450000,
                      36978596.910000, 73957193.820000, 147914387.600000,
                      295828775.300000, 591657550.500000]

        # search nearest zoom of the screen scale
        nearZoomScale = min(LIST_SCALES, key=lambda x:abs(x-scale)) 
        # search the scale for that zoom
        zoom = LIST_ZOOM[LIST_SCALES.index(nearZoomScale)] - 1

        return zoom, nearZoomScale

    def pixelToLatLon(self,tileSize,zoom):
        "Convert tile width px to deg based on"
        "google's tile of 256px"

        # (largura raster / 256 pixels) * (360/2^zoom)
        larg_graus = (float(tileSize)/256.0) * (360.0 / 2.0**float(zoom)) 
    
        return larg_graus

    def rasterScale(self, tileSize, rasterSize):
        "Relation b/ raster size in px and deg. Input rasterSize in deg."

        scale =  float(rasterSize) / float(tileSize)

        return scale
    
    
class ArcGis(object):
    "ARCPY METHODS"

    def __init__ (self):

        self.mxd = arcpy.mapping.MapDocument('CURRENT')
        self.df = arcpy.mapping.ListDataFrames(self.mxd, "Layers") [0]
        self.newlayer = ''
        
    def refresh(self):

        arcpy.RefreshActiveView()
        arcpy.RefreshTOC()

    def rasterExtent(self, raster):

        rasterDesc = arcpy.Describe(raster)
        xmin = rasterDesc.extent.XMin
        ymin = rasterDesc.extent.YMin
        xmax = rasterDesc.extent.XMax
        ymax = rasterDesc.extent.YMax

        return xmin, ymin, xmax, ymax

    def mapExtent(self):

        # 4326 => tell me lat/lon in WGS84/Geograph.
        extent = self.df.extent.projectAs(arcpy.SpatialReference(4326))
        xmin = extent.XMin
        xmax = extent.XMax
        ymin = extent.YMin
        ymax = extent.YMax

        return xmin, ymin, xmax, ymax
        
    def rasterCenter(self, rasterExtent):
        "rasterExtent[xmin, ymin, xmax, ymax]"

        rasterCenterX = (rasterExtent[0] + rasterExtent[2] ) / 2.0 
        rasterCenterY = (rasterExtent[1] + rasterExtent[3] ) / 2.0

        return rasterCenterX, rasterCenterY

    def rasterWidth(self, raster):

        ex = self.rasterExtent(raster)

        return ex[2] - ex[0]
                
    def userMap(self):
        "Get user screen center and proj"

        ex = self.mapExtent()
        mapCenterLon = (ex[0] + ex[2]) / 2.0 
        mapCenterLat = (ex[1] + ex[3]) / 2.0
        mapScale = self.df.scale
        proj = self.df.spatialReference.name

        return mapCenterLat, mapCenterLon, mapScale, proj
    
    def defineWebMercator(self, raster):

        GOOGLE_PROJ = arcpy.SpatialReference(3857)
        arcpy.DefineProjection_management(raster, GOOGLE_PROJ)

    def projectWebMercator(self, raster, rasterLocal, proj):
        "Project raster from WebMercator to whatever"

        IN_COOR_SYSTEM = arcpy.SpatialReference(3857)
        RESSAMPLING_TYPE = "BILINEAR"
        CELL_SIZE = ""
        REGISTRATION_POINT = ""
        TRANSFORMATION = ""

        out_coor_system = arcpy.SpatialReference(proj)
        outRasterFileName = 'raster_projected_' + str(randrange(1000000)) + \
                            '.tif'
        webMercatorRaster = rasterLocal + outRasterFileName

        arcpy.ProjectRaster_management(raster, webMercatorRaster,
                                       out_coor_system , RESSAMPLING_TYPE,
                                       CELL_SIZE, TRANSFORMATION,
                                       REGISTRATION_POINT, IN_COOR_SYSTEM)
        return webMercatorRaster

    def rasterRescale(self, raster, rasterLocal, scale):

        scale_str = str(scale)
        outRasterFileName = 'raster_rescaled_' + str(randrange(1000000)) + \
                            '.tif'
        scaledRaster = rasterLocal + outRasterFileName
        arcpy.Rescale_management(raster, scaledRaster, scale_str, scale_str)

        return scaledRaster

    def rasterMove(self, raster, rasterLocal, mapCenterLat, mapCenterLon,
                   rasterCenter):
        "(raster address, folder address, double, double, array[xmin, ymin, \
        xmax, ymax])"

        outRasterFileName = 'raster_' + str(randrange(1000000)) + '.tif'
        movedScaledRaster = rasterLocal + outRasterFileName

        #colocar uma inteligência aqui para ver ser lat/lon são negativos

        x_relativo = float(mapCenterLon) - (rasterCenter[0])
        y_relativo = float(mapCenterLat) - (rasterCenter[1])
        
        arcpy.Shift_management(raster, movedScaledRaster, x_relativo,
                               y_relativo)

        return movedScaledRaster

    def warpRaster(self, raster, rasterLocal, rasterExtent, targetExtent):
        "distort raster so it fits"

        TRANSF_TYPE = "POLYORDER0"
        RESSAMP_TYPE = "BILINEAR"

        # minimum of 3 pts for warp: just take center point and go.
        rasterCenterX = (rasterExtent[0] + rasterExtent[2]) / 2
        rasterCenterY = (rasterExtent[1] + rasterExtent[3]) / 2
        targetCenterX =  (targetExtent[0] + targetExtent[2]) / 2
        targetCenterY =  (targetExtent[1] + targetExtent[3]) / 2

        srcPts = "\"'" + str(rasterExtent[0]) + " " + str(rasterExtent[1]) +\
                 "';'" + str(rasterExtent[2]) + " " + str(rasterExtent[3]) +\
                 "';'" + str(rasterCenterX) + " " + str(rasterCenterY) + "'\""
        tgtPts = "\"'" + str(targetExtent[0]) + " " + str(targetExtent[1]) +\
                 "';'" + str(targetExtent[2]) + " " + str(targetExtent[3]) + \
                 "';'" + str(targetCenterX) + " " + str(targetCenterY) + "'\""

        outRasterFileName = 'raster_warped_' + str(randrange(1000000)) + \
                            '.tif'
        scaledRaster = rasterLocal + outRasterFileName

        arcpy.AddError('%s %s %s %s %s %s' % (raster,srcPts,tgtPts,
                                              outRasterFileName,TRANSF_TYPE,
                                              RESSAMP_TYPE))

        arcpy.Warp_management(raster, srcPts, tgtPts, outRasterFileName,
                              TRANSF_TYPE, RESSAMP_TYPE)

        return outRasterFileName
        
    def createLayer(self, layer):

        layerBaseName = os.path.basename(layer)
        arcpy.MakeRasterLayer_management(layer, layerBaseName)

        return layerBaseName

    def addLayer(self, layer, position):

        newlayer = arcpy.mapping.Layer(layer)
        arcpy.mapping.AddLayer(self.df, newlayer, position)


class Web(object):
    
    def google(self, mapCenterLat, mapCenterLon, zoomLevel,
               largRaster_px, altRaster_px, GEMapType, rasterLocal):

        url = 'http://maps.googleapis.com/maps/api/staticmap?center=%s,%s& \
                zoom=%s&size=%sx%s&sensor=false&maptype=%s' % \
              (mapCenterLat,mapCenterLon,zoomLevel,largRaster_px,
               altRaster_px,GEMapType)

        try:
            raster = urllib2.urlopen(url)
        except Exception as e:
            print e.message
            arcpy.AddError("Ih... Deu algum problema para acessar o Google"\
                           "Maps. Tente executar o script novamente.\
                            Se persistir, verifique a conexão com a\
                            internet.")
            sys.exit()

        outRasterFileName = 'raster_google_' + str(randrange(1000000)) +\
                            '.jpg'
        outRaster = rasterLocal + outRasterFileName

        GERasterLocal = open(outRaster, 'wb')
        GERasterLocal.write(raster.read())
        GERasterLocal.close()

        return url, GERasterLocal.name

    
class Tools(object):

    def __init__(self, verbose):
        self.verbose = verbose

    def msg(self, msg):
        if self.verbose == "true":
            arcpy.AddWarning("* [ Add Google ] " + msg)

    def error(self, msg):
        arcpy.AddError("* [ Add Google ] ERROR: " + msg)

    def limpeza(self, folder):
        "temp folder cleanup"
                
        self.msg("Tentando limpar a bagunça que eu criei (pasta Temp)...")
        try:
           shutil.rmtree(folder)
        except Exception as f:
            print f.message
            self.error("Não consegui apagar o lixo. Provavelmente o "\
                       "ArcMap.exe tá empacando. Você vai ter que apagar "\
                       "no braço depois de fechar o ArcMap! Se mesmo assim, "\
                       "não conseguir apagar, dê um Ctrl + Alt + Del e "\
                       "manda encerrar o ArcMap.exe.")
            self.error(f.message)

    def memClear(self):
        "vars cleanup"
        for uniquevar in [var for var in globals().copy() if var[0] != "_"
                          and var != 'clearall']:
            del globals()[uniquevar]

# -------------------------------------------------------------------------

# get params
GEMapType = str(sys.argv[1]) 
position = str(sys.argv[2])
apagarTemp = str(sys.argv[3])
verbose = str(sys.argv[4])

#constants
TILE_SIZE = 640

# create objects
coord = GlobalMercator(TILE_SIZE) 
gis = ArcGis()
web = Web()
tools = Tools(verbose)

# verify map
if gis.df.spatialReference.name == '':
    tools.error("O mapa não tem uma referência espacial. " \
             "Adicione ao menos um shapefile ao mapa para " \
             "defini-la ou defina manualmente nas propriedades.")

    tools.error("Abortado.")
    sys.exit()

if gis.df.displayUnits != 'DecimalDegrees':
    tools.error("O mapa não parece estar em graus decimais. "\
                "Poderá haver distorções!")

if gis.df.spatialReference.name != 'GCS_WGS_1984':
    tools.error("O mapa não parece estar em 'GCS_WGS_1984'. "\
                "Poderá haver distorções!")

# reads screen position and check if ok
mapCenterLat = gis.userMap()[0] 
mapCenterLon = gis.userMap()[1]
mapProj = gis.userMap()[2]
isOk = coord.domain(mapCenterLat,mapCenterLon)
if isOk:
    tools.error(isOk)
    sys.exit()
else:
    tools.msg("Limites de Lat/Lon OK")    

# folder basename (get file.mxd and remove .mxd)
try:
    fileNameEnd = os.path.basename(gis.mxd.filePath).index('.')
    fileName = os.path.basename(gis.mxd.filePath)[:fileNameEnd] # remove .mxd 
except Exception as f:
    print f.message
    tools.msg("Salve o arquivo primeiro.")
    self.error(f.message)
    
fdTemp = os.path.dirname(gis.mxd.filePath) + '\\' + fileName + '_Temp\\'
fdRaster = os.path.dirname(gis.mxd.filePath) + '\\' + fileName + '_GE\\'

# folders creation
try:
    os.stat(fdTemp)
except:
    os.mkdir(fdTemp)
try:
    os.stat(fdRaster)
except:
    os.mkdir(fdRaster)

# find out Zoom and change screen scale to equiv.
scale = gis.userMap()[2]
zoomScale = coord.zoomScale(scale) # (zoom,scale)
gis.df.scale = zoomScale[1] # fixa o df.scale
gis.refresh()
tools.msg('Centro do mapa: %s %s / Zoom: %s / Escala Equiv: 1:%s' % \
          (mapCenterLat,mapCenterLon,zoomScale[0],zoomScale[1]))

# get google map tile
rasterGoogle = web.google(str(mapCenterLat), str(mapCenterLon),
                          str(zoomScale[0]), TILE_SIZE, TILE_SIZE,
                          GEMapType, fdTemp)
tools.msg("Download: " + rasterGoogle[0])

# define raster projection
gis.defineWebMercator(rasterGoogle[1])
tools.msg("Definida projecao 3857 para o raster")

# project raster (ATUALMENTE PARA WGS84 FORÇADO)
rasterProj = gis.projectWebMercator(rasterGoogle[1], fdTemp, 4326)
tools.msg("Raster projetado: %s" % rasterProj)

# rescale raster
rasterWidth = coord.pixelToLatLon(TILE_SIZE,zoomScale[0])
rasterWidth_inMap = gis.rasterWidth(rasterProj)
rasterScale = coord.rasterScale(rasterWidth_inMap, rasterWidth)
rasterRedim = gis.rasterRescale(rasterProj, fdTemp, rasterScale)
tools.msg("Raster Width: %s" % rasterWidth)
tools.msg("Fator de redimensionamento (x,y): %s" % rasterScale)

# raster props
rasterExtent = gis.rasterExtent(rasterRedim)
rasterCenter = gis.rasterCenter(rasterExtent)

###TESTE COM WARP
##GOOGLE_PROJ = arcpy.SpatialReference(4326)
##arcpy.DefineProjection_management(rasterGoogle[1], GOOGLE_PROJ)
##
##newLayer = gis.createLayer(rasterGoogle[1])
##gis.addLayer(newLayer, position)
##tools.msg("Adicionada camada: %s" % newLayer)
##
##mapExtent = gis.mapExtent()
##rasterExtent = gis.rasterExtent(newLayer)
##rasterRect = gis.warpRaster(newLayer, fdTemp, rasterExtent, mapExtent)
##tools.error(str(rasterRect))

# move raster
rasterReposic = gis.rasterMove(rasterRedim, fdRaster, mapCenterLat,
                               mapCenterLon, rasterCenter)
tools.msg("Raster Movido: %s" % rasterReposic)

# create and add layer
newLayer = gis.createLayer(rasterReposic)
gis.addLayer(newLayer, position)
tools.msg("Adicionada camada: %s" % newLayer)

# refresh
gis.refresh()
tools.msg("Atualizando...")

# cleanup folders
if apagarTemp == "true":
    tools.limpeza(fdTemp)

# cleanup vars
tools.memClear()